import json
import pathlib
import sys
import time
import urllib.request
import logging
from datetime import datetime
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QHBoxLayout, 
                             QWidget, QTextEdit, QPushButton, QLabel, QLineEdit, 
                             QCheckBox, QProgressBar, QMessageBox, QScrollArea, QSplitter)
from PyQt5.QtCore import QThread, pyqtSignal, Qt
from PyQt5.QtGui import QPixmap, QFont, QIcon

OUTPUT_DIR = pathlib.Path(__file__).parent / "output"
OUTPUT_DIR.mkdir(exist_ok=True, parents=True)

# 设置日志
LOG_DIR = pathlib.Path(__file__).parent / "logs"
LOG_DIR.mkdir(exist_ok=True, parents=True)

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(LOG_DIR / f"comfyui_client_{datetime.now().strftime('%Y%m%d')}.log", encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class ImageGenerationWorker(QThread):
    """图像生成工作线程，防止UI冻结"""
    finished = pyqtSignal(str, float)  # filename, elapsed_time
    error = pyqtSignal(str)  # error_message
    progress = pyqtSignal(str)  # status_message
    log_signal = pyqtSignal(str)  # log_message
    
    def __init__(self, workspace, prompt, dev_mode=False):
        super().__init__()
        self.workspace = workspace
        self.prompt = prompt
        self.dev_mode = dev_mode
    
    def run(self):
        session_id = datetime.now().strftime('%H:%M:%S')
        try:
            self.log_signal.emit(f"[{session_id}] 开始图像生成任务")
            self.log_signal.emit(f"[{session_id}] 工作空间: {self.workspace}")
            self.log_signal.emit(f"[{session_id}] 开发模式: {'是' if self.dev_mode else '否'}")
            self.log_signal.emit(f"[{session_id}] 提示词: {self.prompt}")
            
            logger.info(f"开始图像生成 - 工作空间: {self.workspace}, 开发模式: {self.dev_mode}")
            logger.info(f"提示词: {self.prompt}")
            
            url = f"https://{self.workspace}--example-comfyapp-comfyui-api{'-dev' if self.dev_mode else ''}.modal.run/"
            self.log_signal.emit(f"[{session_id}] 构建请求URL: {url}")
            
            data = json.dumps({"prompt": self.prompt}).encode("utf-8")
            self.log_signal.emit(f"[{session_id}] 准备请求数据，大小: {len(data)} 字节")
            
            self.progress.emit("正在发送请求到服务器...")
            self.log_signal.emit(f"[{session_id}] 发送HTTP请求到服务器")
            
            start_time = time.time()
            req = urllib.request.Request(
                url, data=data, headers={"Content-Type": "application/json"}
            )
            
            logger.info(f"发送请求到: {url}")
            
            self.progress.emit("服务器正在生成图像，请稍候...")
            self.log_signal.emit(f"[{session_id}] 等待服务器响应...")
            
            with urllib.request.urlopen(req) as response:
                self.log_signal.emit(f"[{session_id}] 收到服务器响应，状态码: {response.status}")
                
                if response.status != 200:
                    error_msg = f"HTTP错误: {response.status}"
                    logger.error(error_msg)
                    self.log_signal.emit(f"[{session_id}] 错误: {error_msg}")
                    self.error.emit(error_msg)
                    return
                
                self.log_signal.emit(f"[{session_id}] 开始下载图像数据...")
                image_data = response.read()
                self.log_signal.emit(f"[{session_id}] 图像数据下载完成，大小: {len(image_data)} 字节")
                
                elapsed = round(time.time() - start_time, 1)
                filename = OUTPUT_DIR / f"{slugify(self.prompt)}.png"
                
                self.log_signal.emit(f"[{session_id}] 保存图像到: {filename}")
                filename.write_bytes(image_data)
                
                logger.info(f"图像生成成功 - 耗时: {elapsed}秒, 保存到: {filename}")
                self.log_signal.emit(f"[{session_id}] 图像生成完成！耗时: {elapsed}秒")
                
                self.finished.emit(str(filename), elapsed)
                
        except urllib.error.HTTPError as e:
            error_msg = ""
            if e.code == 404:
                error_msg = f"未找到工作流API: {url}"
            else:
                error_msg = f"HTTP错误 {e.code}: {e.reason}"
            
            logger.error(f"HTTP错误: {error_msg}")
            self.log_signal.emit(f"[{session_id}] HTTP错误: {error_msg}")
            self.error.emit(error_msg)
            
        except Exception as e:
            error_msg = f"未知错误: {str(e)}"
            logger.error(f"生成图像时发生错误: {error_msg}")
            self.log_signal.emit(f"[{session_id}] 发生错误: {error_msg}")
            self.error.emit(error_msg)


class ComfyUIClient(QMainWindow):
    def __init__(self):
        super().__init__()
        self.worker = None
        self.load_styles()
        self.init_ui()
    
    def load_styles(self):
        """加载QSS样式文件"""
        try:
            style_file = pathlib.Path(__file__).parent / "styles.qss"
            if style_file.exists():
                with open(style_file, 'r', encoding='utf-8') as f:
                    self.setStyleSheet(f.read())
                logger.info("样式文件加载成功")
            else:
                logger.warning("样式文件不存在，使用默认样式")
        except Exception as e:
            logger.error(f"加载样式文件失败: {e}")
            # 如果加载失败，使用内置的基础样式
            self.setStyleSheet(self.get_fallback_style())
        
    def init_ui(self):
        self.setWindowTitle("ComfyUI 客户端 - AI图像生成器")
        self.setGeometry(100, 100, 1200, 800)
        self.setMinimumSize(1000, 600)
        
        # 中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 主布局
        main_layout = QHBoxLayout(central_widget)
        main_layout.setSpacing(10)
        main_layout.setContentsMargins(10, 10, 10, 10)
        
        # 左侧控制面板
        left_panel = QWidget()
        left_panel.setObjectName("leftPanel")
        left_panel.setMaximumWidth(380)
        left_panel.setMinimumWidth(350)
        left_layout = QVBoxLayout(left_panel)
        left_layout.setSpacing(12)
        left_layout.setContentsMargins(15, 15, 15, 15)
        
        # 标题
        title = QLabel("🎨 ComfyUI 图像生成器")
        title.setObjectName("titleLabel")
        title.setAlignment(Qt.AlignCenter)
        left_layout.addWidget(title)
        
        # 工作空间输入区域
        workspace_label = QLabel("🌐 Modal 工作空间:")
        left_layout.addWidget(workspace_label)
        self.workspace_input = QLineEdit()
        self.workspace_input.setPlaceholderText("例如: modal-labs-examples")
        self.workspace_input.setToolTip("输入您的Modal工作空间名称")
        left_layout.addWidget(self.workspace_input)
        
        # 开发模式复选框
        self.dev_checkbox = QCheckBox("🔧 开发模式")
        self.dev_checkbox.setToolTip("启用开发模式以连接到开发服务器")
        left_layout.addWidget(self.dev_checkbox)
        
        # 提示词输入区域
        prompt_label = QLabel("✨ 提示词:")
        left_layout.addWidget(prompt_label)
        self.prompt_input = QTextEdit()
        self.prompt_input.setMaximumHeight(120)
        self.prompt_input.setPlaceholderText("在此输入您的图像生成提示词...\n例如: 一只可爱的小猫坐在彩虹上，卡通风格，高质量")
        self.prompt_input.setToolTip("输入详细的图像描述以获得更好的生成效果")
        left_layout.addWidget(self.prompt_input)
        
        # 生成按钮
        self.generate_btn = QPushButton("🚀 生成图像")
        self.generate_btn.setToolTip("点击开始生成图像")
        self.generate_btn.clicked.connect(self.generate_image)
        left_layout.addWidget(self.generate_btn)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        self.progress_bar.setTextVisible(True)
        left_layout.addWidget(self.progress_bar)
        
        # 状态标签
        self.status_label = QLabel("✅ 就绪")
        self.status_label.setObjectName("statusLabel")
        self.status_label.setWordWrap(True)
        left_layout.addWidget(self.status_label)
        
        left_layout.addStretch()
        
        # 右侧面板 - 使用分割器分为图像显示和日志
        right_panel = QSplitter(Qt.Vertical)
        
        # 图像显示区域
        image_widget = QWidget()
        image_widget.setObjectName("imageWidget")
        image_layout = QVBoxLayout(image_widget)
        image_layout.setContentsMargins(15, 15, 15, 15)
        
        image_title = QLabel("🖼️ 生成的图像:")
        image_layout.addWidget(image_title)
        
        self.image_scroll = QScrollArea()
        self.image_label = QLabel()
        self.image_label.setObjectName("imageDisplayLabel")
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setText("🎨 生成的图像将在此显示\n\n点击上方的'生成图像'按钮开始创作")
        
        self.image_scroll.setWidget(self.image_label)
        self.image_scroll.setWidgetResizable(True)
        image_layout.addWidget(self.image_scroll)
        
        # 日志显示区域
        log_widget = QWidget()
        log_widget.setObjectName("logWidget")
        log_layout = QVBoxLayout(log_widget)
        log_layout.setContentsMargins(15, 15, 15, 15)
        
        log_title = QLabel("📋 详细日志:")
        log_layout.addWidget(log_title)
        
        self.log_display = QTextEdit()
        self.log_display.setObjectName("logDisplay")
        self.log_display.setMaximumHeight(200)
        self.log_display.setReadOnly(True)
        log_layout.addWidget(self.log_display)
        
        # 添加到分割器
        right_panel.addWidget(image_widget)
        right_panel.addWidget(log_widget)
        right_panel.setSizes([500, 200])  # 设置初始大小比例
        
        # 添加面板到主布局
        main_layout.addWidget(left_panel)
        main_layout.addWidget(right_panel, 1)
        
        # 设置默认工作空间
        self.workspace_input.setText("modal-labs-examples")
        
        # 初始化日志
        self.add_log("🎉 程序启动完成，等待用户操作...")
        logger.info("ComfyUI客户端启动完成")
        
    def add_log(self, message):
        """添加日志到显示区域，支持颜色编码"""
        timestamp = datetime.now().strftime('%H:%M:%S')
        
        # 根据消息内容设置颜色
        if "错误" in message or "失败" in message or "Error" in message:
            color = "#ff6b6b"  # 红色
            icon = "❌"
        elif "警告" in message or "Warning" in message:
            color = "#ffa726"  # 橙色
            icon = "⚠️"
        elif "完成" in message or "成功" in message or "Success" in message:
            color = "#4caf50"  # 绿色
            icon = "✅"
        elif "开始" in message or "启动" in message:
            color = "#2196f3"  # 蓝色
            icon = "🚀"
        elif "状态更新" in message:
            color = "#9c27b0"  # 紫色
            icon = "📊"
        else:
            color = "#00ff00"  # 默认绿色
            icon = "ℹ️"
        
        # 格式化日志消息
        formatted_message = f'<span style="color: {color};">{icon} [{timestamp}] {message}</span>'
        self.log_display.append(formatted_message)
        
        # 自动滚动到底部
        scrollbar = self.log_display.verticalScrollBar()
        scrollbar.setValue(scrollbar.maximum())
    
    def generate_image(self):
        workspace = self.workspace_input.text().strip()
        prompt = self.prompt_input.toPlainText().strip()
        
        self.add_log("开始验证输入参数...")
        
        if not workspace:
            error_msg = "请输入Modal工作空间名称"
            self.add_log(f"验证失败: {error_msg}")
            QMessageBox.warning(self, "警告", error_msg)
            return
            
        if not prompt:
            error_msg = "请输入图像生成提示词"
            self.add_log(f"验证失败: {error_msg}")
            QMessageBox.warning(self, "警告", error_msg)
            return
        
        self.add_log("输入参数验证通过")
        self.add_log(f"工作空间: {workspace}")
        self.add_log(f"开发模式: {'是' if self.dev_checkbox.isChecked() else '否'}")
        self.add_log(f"提示词长度: {len(prompt)} 字符")
        
        # 在生成期间禁用UI
        self.generate_btn.setEnabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setRange(0, 0)  # 不确定进度
        self.status_label.setText("正在初始化...")
        
        self.add_log("UI已锁定，开始启动工作线程...")
        
        # 启动工作线程
        self.worker = ImageGenerationWorker(workspace, prompt, self.dev_checkbox.isChecked())
        self.worker.finished.connect(self.on_generation_finished)
        self.worker.error.connect(self.on_generation_error)
        self.worker.progress.connect(self.on_progress_update)
        self.worker.log_signal.connect(self.add_log)
        self.worker.start()
        
        self.add_log("工作线程已启动")
    
    def on_generation_finished(self, filename, elapsed_time):
        # 重新启用UI
        self.generate_btn.setEnabled(True)
        self.progress_bar.setVisible(False)
        success_msg = f"图像生成完成！耗时: {elapsed_time}秒 - 保存到: {filename}"
        self.status_label.setText(success_msg)
        
        self.add_log("开始加载生成的图像...")
        
        # 加载并显示图像
        pixmap = QPixmap(filename)
        if not pixmap.isNull():
            # 缩放图像以适应显示区域，保持宽高比
            scaled_pixmap = pixmap.scaled(600, 600, Qt.KeepAspectRatio, Qt.SmoothTransformation)
            self.image_label.setPixmap(scaled_pixmap)
            self.image_label.setText("")
            
            # 获取图像信息
            image_size = f"{pixmap.width()}x{pixmap.height()}"
            file_size = pathlib.Path(filename).stat().st_size
            file_size_mb = round(file_size / 1024 / 1024, 2)
            
            self.add_log(f"图像加载成功 - 尺寸: {image_size}, 文件大小: {file_size_mb}MB")
            logger.info(f"图像显示成功 - 文件: {filename}, 尺寸: {image_size}")
        else:
            error_msg = "无法加载生成的图像"
            self.image_label.setText(error_msg)
            self.add_log(f"错误: {error_msg}")
            logger.error(f"无法加载图像文件: {filename}")
        
        self.add_log("图像生成流程完成")
    
    def on_generation_error(self, error_message):
        # 重新启用UI
        self.generate_btn.setEnabled(True)
        self.progress_bar.setVisible(False)
        self.status_label.setText(f"错误: {error_message}")
        
        self.add_log(f"生成失败: {error_message}")
        QMessageBox.critical(self, "生成错误", error_message)
    
    def on_progress_update(self, message):
        self.status_label.setText(message)
        # 进度更新也记录到日志（但不重复记录相同信息）
        if not message.startswith("["):  # 避免重复记录已经有时间戳的消息
            self.add_log(f"状态更新: {message}")
    
    def get_fallback_style(self):
        """备用样式，当QSS文件加载失败时使用"""
        return """
        QMainWindow {
            background-color: #f5f5f5;
            font-family: "Microsoft YaHei", Arial, sans-serif;
        }
        QPushButton {
            background: qlineargradient(x1: 0, y1: 0, x2: 0, y2: 1,
                                        stop: 0 #667eea, stop: 1 #764ba2);
            color: white;
            border: none;
            border-radius: 8px;
            padding: 12px 24px;
            font-size: 14px;
            font-weight: 600;
        }
        QPushButton:hover {
            background: qlineargradient(x1: 0, y1: 0, x2: 0, y2: 1,
                                        stop: 0 #5a6fd8, stop: 1 #6a4190);
        }
        QPushButton:disabled {
            background-color: #adb5bd;
        }
        QLineEdit, QTextEdit {
            background-color: white;
            border: 2px solid #e9ecef;
            border-radius: 6px;
            padding: 8px;
            font-size: 13px;
        }
        QLineEdit:focus, QTextEdit:focus {
            border-color: #667eea;
        }
        """


def slugify(s: str) -> str:
    return s.lower().replace(" ", "-").replace(".", "-").replace("/", "-")[:32]


def main():
    logger.info("启动ComfyUI客户端应用程序")
    
    # 创建应用程序实例
    app = QApplication(sys.argv)
    app.setApplicationName("ComfyUI客户端")
    app.setApplicationVersion("2.0")
    app.setApplicationDisplayName("ComfyUI AI图像生成器")
    
    # 设置应用程序图标（如果有的话）
    try:
        # 可以在这里设置应用程序图标
        # app.setWindowIcon(QIcon("icon.png"))
        pass
    except:
        pass
    
    # 创建主窗口
    window = ComfyUIClient()
    window.show()
    
    logger.info("应用程序界面已显示，样式已加载")
    
    # 运行应用程序
    try:
        result = app.exec_()
        logger.info("应用程序正常退出")
        return result
    except Exception as e:
        logger.error(f"应用程序异常退出: {e}")
        return 1


if __name__ == "__main__":
    main()
